import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import cont2discrete
from MPC_InputAugmentMatrix_SS_U import F2_InputAugmentMatrix_SS_U
from MPC_Matrices_PM import F4_MPC_Matrices_PM
from MPC_controller_withConstriants import F7_MPC_Controller_withConstraints
from MPC_Matrices_Constraints import F6_MPC_Matrices_Constraints
import socket
import pickle
import json
import time
class MPC():
    def __init__(self) -> None:
        # 定义预测区间
        N_P = 20

        ## 定义系统参数
        # 定义K-T常数
        K=3.1429
        T=18.348

        ## 系统定义
        # 构建系统矩阵A，n x n
        self.A = np.array([[0, 1],
                [0, 1/T]])
        # 计算A矩阵维度
        n = np.size(self.A,axis=0)
        # print(n)
        # 构建输入矩阵B，n x p
        self.B = np.array([[0], 
                [K/T]
                ])
        self.p = np.size(self.B, axis=1)

        ## 系统离散
        Ts = 0.1
        # 连续系统转离散系统 零阶保持法离散化公式
        sys_d = cont2discrete((self.A, self.B, np.eye(n), np.zeros((n, self.p))), Ts)

        # 提取离散系统A矩阵
        self.A = sys_d[0]
        # 提取离散系统B矩阵
        self.B = sys_d[1]

        ## 权重设计
        # 设计状态权重系数矩阵，n x n
        self.Q = np.array( [
            [1, 0],
            [0, 1]
        ])

        # 设计终值权重系统矩阵，n x n
        self.S = np.array([
            [1, 0],
            [0, 0]
        ])

        # 设计输入权重系数矩阵，p x p
        self.R = np.array([
            [0.1]
        ])
        
        ## 系统参考值
        # 系统状态参考值
        self.xd = np.array([
            [90],
            [0]
        ])
        # 构建目标转移矩阵
        AD = np.eye(n)
        # 定义预测区间
        self.N_P = 20
        
        # 计算目标输入
        # 目标输入的计算: ud = B \ ((I - A) * xd)
        # 相当于在 Python 中求解 B * ud = (I - A) * xd
        self.ud = np.linalg.pinv(self.B) @ (np.eye(n) - self.A) @ self.xd
        # print(ud)

        ## 初始化系统状态
        self.x0 = np.array([
            [0],
            [0]
        ])
           
        # 初始化系统输入
        u0 = np.array([[0]])
        u = u0

        ## 系统约束定义
        # 输入下限
        self.u_low = np.array([[-30]])
        # 输入上限
        self.u_high = np.array([[30]])

        # 状态下限
        self.x_low = np.array([
            [-180],
            [-15.0]
        ])
        # 状态上限
        self.x_high = np.array([
            [180],
            [15.0]
        ])
        # 增广状态下限
        self.xa_low = np.vstack([
            self.x_low,
            [[float('-inf')],
            [float('-inf')]]
        ])
        # 增广状态上限
        self.xa_high = np.vstack([
            self.x_high,
            [[float('inf')],
            [float('inf')]]
        ])
        
    def solver(self):
        x = self.x0
        # 初始化增广状态矩阵
        xa = np.vstack(
            [x, self.xd]
        )
        # 调用模块F2 计算系统增广矩阵 Aa, Ba, Qa, Sa, R 以及目标输入ud
        # 对于非定常xd每次都要计算，应该都是在线计算的了
        Aa, Ba, Qa, Sa, R, ud = F2_InputAugmentMatrix_SS_U(self.A, self.B, self.Q, self.R, self.S, self.xd)

        # 调用模块F4 计算二次规划需要用到的矩阵
        Phi,Gamma,Omega,Psi,F,H = F4_MPC_Matrices_PM(Aa, Ba, Qa, R, Sa, self.N_P)

        # 调用模块F6 计算含约束二次规划需要用到的矩阵
        M,Beta_bar,b = F6_MPC_Matrices_Constraints(self.xa_low, self.xa_high, self.u_low, self.u_high, self.N_P, Phi, Gamma)
        
        # 调用模块F7 计算系统系统控制（输入增量）
        delta_U, delta_u = F7_MPC_Controller_withConstraints(xa,F,H,M,Beta_bar,b, self.p)

        # 根据输入增量计算系统输入
        u = delta_u + ud
        return u


def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    mpc_state_dict = dict()
    mpc_state_dict = recv_obser['mpc']

    reward = recv_obser['reward']
    terminated = recv_obser['terminated']
    truncated = recv_obser['truncated']
    return mpc_state_dict, reward, terminated, truncated


def encode_data(action, reset_flag=1):
    rudder = float(action) # add?

    if rudder >= 30.0:
        rudder = 30.0
    elif rudder <= -30.0:
        rudder = -30.0
    
    origin_data = {'boatname':'SLM7001',
                   'restart': reset_flag,
                   'rudl': rudder,
                   'rudr': rudder,
                   'rspl': 1000,
                   'rspr': 1000,
                   'subSystem': "control"
                   }

    data = json.dumps(origin_data, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8')


if __name__=='__main__':
    # init socket
    # network socket
    REMOTE_HOST = '127.0.0.1'
    REMOTE_PORT = 10906
    tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_socket.connect((REMOTE_HOST, REMOTE_PORT))
    
    controller_mpc = MPC()
    
    record = dict()
    # 定义x_history 零矩阵，用于存储系统状态结果，维度n x k_step
    record['x1_history'] = list()
    record['x2_history'] = list()
    # 定义u_history 零矩阵，用于存储系统输入结果，维度p x k_step
    record['u_history'] = list()
    u = 0
    
    mean_time = 0
    count = 0
    while(1):
        count += 1
        
        ## 系统参考值
        action_data = encode_data(u, reset_flag=1)

        # print(data)
        tcp_socket.send(action_data)
        # Initialize the environment and get its state
        info = tcp_socket.recv(1024)
        mpc_state, _, _, _= decode_data(info)
        
        # 系统状态参考值
        controller_mpc.xd = np.array([
            [mpc_state['psi_d']],
            [0.0]
        ])
        
        controller_mpc.x0 = np.array([
            [mpc_state['psi']],
            [mpc_state['dpsi']]
        ])
        
        # 计算平均单步运行时间
        time_start = time.time()
        
        u = controller_mpc.solver()
        
        time_end = time.time()
        consume_time = (time_end - time_start) * 1000.0
        
        mean_time = mean_time + consume_time
        
        print(f"consume time is {mean_time/count} ms")
